[CDK]Step Functions StateMachineの実行履歴を自動でAthenaに連携する

cdk init --language typescript

npm install @aws-cdk/aws-glue-alpha @aws-cdk/aws-scheduler-alpha @aws-cdk/aws-scheduler-targets-alpha

// 必要なモジュールをインポート
import * as cdk from "aws-cdk-lib";
import { aws_lambda_nodejs, TimeZone } from "aws-cdk-lib";
import * as lambda from "aws-cdk-lib/aws-lambda";
import * as s3 from "aws-cdk-lib/aws-s3";
import * as athena from "aws-cdk-lib/aws-athena";
import * as iam from "aws-cdk-lib/aws-iam";
import * as glue from "@aws-cdk/aws-glue-alpha";
import * as sfn from "aws-cdk-lib/aws-stepfunctions";
import * as path from "path";
import { Construct } from "constructs";
import * as aws_scheduler_alpha from "@aws-cdk/aws-scheduler-alpha";
import * as aws_scheduler_targets_alpha from "@aws-cdk/aws-scheduler-targets-alpha";

// CDKスタッククラスを定義
export class CdkStepfunctionsMetricStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    // メトリクスとAthenaクエリ結果を保存するS3バケットを作成
    const bucket = new s3.Bucket(this, "StateMachineMetricsBucket", {
      bucketName: "<your-bucket-name>",
      removalPolicy: cdk.RemovalPolicy.DESTROY,
      autoDeleteObjects: true,
    });

    // Athenaワークグループを作成
    const workgroup = new athena.CfnWorkGroup(
      this,
      "StateMachinesMetricsWorkgroup",
      {
        name: "StateMachinesMetricsWorkgroup",
        recursiveDeleteOption: true,
        workGroupConfiguration: {
          resultConfiguration: {
            outputLocation: `s3://${bucket.bucketName}/athena-results/`,
          },
          publishCloudWatchMetricsEnabled: true,
          enforceWorkGroupConfiguration: true,
          bytesScannedCutoffPerQuery: 1073741824,
        },
      }
    );

    // Glueデータベースを作成
    const database = new glue.Database(this, "StateMachinesMetricsDatabase", {
      databaseName: "statemachines_metrics_db",
    });

    // Glueテーブルを作成（Athenaでクエリ可能）
    const table = new glue.S3Table(this, "StateMachinesMetricsTable", {
      database: database,
      tableName: "statemachines_metrics",
      columns: [
        { name: "name", type: glue.Schema.STRING },
        { name: "date", type: glue.Schema.STRING },
        { name: "starttime", type: glue.Schema.STRING },
        { name: "endtime", type: glue.Schema.STRING },
        { name: "duration", type: glue.Schema.DOUBLE },
        { name: "status", type: glue.Schema.STRING },
      ],
      dataFormat: glue.DataFormat.JSON,
      bucket: bucket,
      s3Prefix: "statemachines_metrics/",
    });

    // Lambda関数用のIAMロールを作成
    const lambdaRole = new iam.Role(this, "StateMachinesMetricsLambdaRole", {
      assumedBy: new iam.ServicePrincipal("lambda.amazonaws.com"),
    });

    // Lambda実行に必要な基本的な権限を追加
    lambdaRole.addManagedPolicy(
      iam.ManagedPolicy.fromAwsManagedPolicyName(
        "service-role/AWSLambdaBasicExecutionRole"
      )
    );

    // 必要なAWSサービスへのアクセス権限を追加
    lambdaRole.addToPolicy(
      new iam.PolicyStatement({
        actions: [
          "athena:StartQueryExecution",
          "states:GetExecutionHistory",
          "states:DescribeExecution",
          "states:ListStateMachines",
          "states:ListExecutions",
        ],
        resources: ["*"],
      })
    );

    // Glue Data Catalogへのアクセス権限を追加
    lambdaRole.addToPolicy(
      new iam.PolicyStatement({
        actions: ["glue:GetTable", "glue:GetPartitions"],
        resources: [
          "arn:aws:glue:*:*:catalog",
          `arn:aws:glue:*:*:database/${database.databaseName}`,
          `arn:aws:glue:*:*:table/${database.databaseName}/${table.tableName}`,
        ],
      })
    );

    // Lambda関数を作成
    const lambdaFunction = new aws_lambda_nodejs.NodejsFunction(
      this,
      "StateMachinesMetricsLambda",
      {
        runtime: lambda.Runtime.NODEJS_20_X,
        handler: "handler",
        entry: path.join(__dirname, "../lambda/index.ts"),
        role: lambdaRole,
        timeout: cdk.Duration.minutes(5),
        environment: {
          BUCKET_NAME: bucket.bucketName,
          ATHENA_WORKGROUP: workgroup.name,
          DATABASE_NAME: database.databaseName,
          TABLE_NAME: table.tableName,
        },
      }
    );

    // Lambda関数にS3バケットへの読み書き権限を付与
    bucket.grantReadWrite(lambdaFunction);

    // Lambda関数を実行するためのターゲットを作成
    const target = new aws_scheduler_targets_alpha.LambdaInvoke(
      lambdaFunction,
      {}
    );

    // Lambda関数を定期的に実行するスケジュールを作成
    new aws_scheduler_alpha.Schedule(this, "Schedule", {
      scheduleName: "invoke-lambda-schedule",
      schedule: aws_scheduler_alpha.ScheduleExpression.cron({
        minute: "0",
        hour: "3",
        day: "*",
        month: "*",
        year: "*",
        timeZone: TimeZone.ASIA_TOKYO,
      }),
      target,
    });

    // 5つのStateMachineを作成
    this.createMultipleStateMachines(5);
  }

  // 複数のStateMachineを作成するメソッド
  private createMultipleStateMachines(count: number) {
    for (let i = 1; i <= count; i++) {
      // ランダムで5-300秒の待機時間を生成
      const waitTime = this.generateRandomWaitTime(5, 300);

      // StateMachineの定義を作成
      const definition = sfn.Chain.start(
        new sfn.Wait(this, `Wait${i}`, {
          time: sfn.WaitTime.duration(cdk.Duration.seconds(waitTime)),
        })
      ).next(new sfn.Succeed(this, `Success${i}`));

      // StateMachineを作成
      const stateMachine = new sfn.StateMachine(this, `StateMachine${i}`, {
        definition,
        stateMachineName: `TestStateMachine${i}`,
        timeout: cdk.Duration.minutes(5),
      });

      // 各StateMachine用のスケジューラーを作成
      this.createSchedulerForStateMachine(stateMachine, i);
    }
  }

  // 指定された範囲でランダムな秒数を生成するメソッド
  private generateRandomWaitTime(min: number, max: number): number {
    return Math.floor(Math.random() * (max - min + 1) + min);
  }

  // StateMachineを実行するEventBridge Schedulerを設定するメソッド
  private createSchedulerForStateMachine(
    stateMachine: sfn.StateMachine,
    index: number
  ) {
    const target = new aws_scheduler_targets_alpha.StepFunctionsStartExecution(
      stateMachine,
      {}
    );

    new aws_scheduler_alpha.Schedule(this, `Schedule-StateMachine${index}`, {
      scheduleName: `invoke-StateMachine${index}-schedule`,
      schedule: aws_scheduler_alpha.ScheduleExpression.cron({
        minute: "0",
        hour: "0",
        day: "*",
        month: "*",
        year: "*",
        timeZone: TimeZone.ASIA_TOKYO,
      }),
      target,
    });
  }
}

// メトリクスとAthenaクエリ結果を保存するS3バケットを作成
const bucket = new s3.Bucket(this, "StateMachineMetricsBucket", {
  bucketName: "<your-bucket-name>",
  removalPolicy: cdk.RemovalPolicy.DESTROY,
  autoDeleteObjects: true,
});

// Athenaワークグループを作成
const workgroup = new athena.CfnWorkGroup(
  this,
  "StateMachinesMetricsWorkgroup",
  {
    name: "StateMachinesMetricsWorkgroup",
    recursiveDeleteOption: true,
    workGroupConfiguration: {
      resultConfiguration: {
        outputLocation: `s3://${bucket.bucketName}/athena-results/`,
      },
      publishCloudWatchMetricsEnabled: true,
      enforceWorkGroupConfiguration: true,
      bytesScannedCutoffPerQuery: 1073741824,
    },
  }
);

// Glueデータベースを作成
const database = new glue.Database(this, "StateMachinesMetricsDatabase", {
  databaseName: "statemachines_metrics_db",
});

// Glueテーブルを作成（Athenaでクエリ可能）
const table = new glue.S3Table(this, "StateMachinesMetricsTable", {
  database: database,
  tableName: "statemachines_metrics",
  columns: [
    { name: "name", type: glue.Schema.STRING },
    { name: "date", type: glue.Schema.STRING },
    { name: "starttime", type: glue.Schema.STRING },
    { name: "endtime", type: glue.Schema.STRING },
    { name: "duration", type: glue.Schema.DOUBLE },
    { name: "status", type: glue.Schema.STRING },
  ],
  dataFormat: glue.DataFormat.JSON,
  bucket: bucket,
  s3Prefix: "statemachines_metrics/",
});

// Lambda関数用のIAMロールを作成
const lambdaRole = new iam.Role(this, "StateMachinesMetricsLambdaRole", {
  assumedBy: new iam.ServicePrincipal("lambda.amazonaws.com"),
});

// Lambda実行に必要な基本的な権限を追加
lambdaRole.addManagedPolicy(
  iam.ManagedPolicy.fromAwsManagedPolicyName(
    "service-role/AWSLambdaBasicExecutionRole"
  )
);

// 必要なAWSサービスへのアクセス権限を追加
lambdaRole.addToPolicy(
  new iam.PolicyStatement({
    actions: [
      "athena:StartQueryExecution",
      "states:GetExecutionHistory",
      "states:DescribeExecution",
      "states:ListStateMachines",
      "states:ListExecutions",
    ],
    resources: ["*"],
  })
);

// Glue Data Catalogへのアクセス権限を追加
lambdaRole.addToPolicy(
  new iam.PolicyStatement({
    actions: ["glue:GetTable", "glue:GetPartitions"],
    resources: [
      "arn:aws:glue:*:*:catalog",
      `arn:aws:glue:*:*:database/${database.databaseName}`,
      `arn:aws:glue:*:*:table/${database.databaseName}/${table.tableName}`,
    ],
  })
);

// Lambda関数を作成
const lambdaFunction = new aws_lambda_nodejs.NodejsFunction(
  this,
  "StateMachinesMetricsLambda",
  {
    runtime: lambda.Runtime.NODEJS_20_X,
    handler: "handler",
    entry: path.join(__dirname, "../lambda/index.ts"),
    role: lambdaRole,
    timeout: cdk.Duration.minutes(5),
    environment: {
      BUCKET_NAME: bucket.bucketName,
      ATHENA_WORKGROUP: workgroup.name,
      DATABASE_NAME: database.databaseName,
      TABLE_NAME: table.tableName,
    },
  }
);

// Lambda関数にS3バケットへの読み書き権限を付与
bucket.grantReadWrite(lambdaFunction);

// Lambda関数を実行するためのターゲットを作成
const target = new aws_scheduler_targets_alpha.LambdaInvoke(
  lambdaFunction,
  {}
);

// Lambda関数を定期的に実行するスケジュールを作成
new aws_scheduler_alpha.Schedule(this, "Schedule", {
  scheduleName: "invoke-lambda-schedule",
  schedule: aws_scheduler_alpha.ScheduleExpression.cron({
    minute: "0",
    hour: "3",
    day: "*",
    month: "*",
    year: "*",
    timeZone: TimeZone.ASIA_TOKYO,
  }),
  target,
});

export class CdkStateMachinesMetricsStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    // 省略

    // 5つのStateMachineを作成
    this.createMultipleStateMachines(5);
  }

  private createMultipleStateMachines(count: number) {
    for (let i = 1; i <= count; i++) {
      // ランダムで5-300秒を取得する
      const waitTime = this.generateRandomWaitTime(5, 300);

      const definition = sfn.Chain.start(
        new sfn.Wait(this, `Wait${i}`, {
          time: sfn.WaitTime.duration(cdk.Duration.seconds(waitTime)),
        })
      ).next(new sfn.Succeed(this, `Success${i}`));

      const stateMachine = new sfn.StateMachine(this, `StateMachine${i}`, {
        definition,
        stateMachineName: `TestStateMachine${i}`,
        timeout: cdk.Duration.minutes(5),
      });

      // Create a scheduler for each StateMachine
      this.createSchedulerForStateMachine(stateMachine, i);
    }
  }

  // 指定された範囲でランダムな秒数を取得する関数
  private generateRandomWaitTime(min: number, max: number): number {
    return Math.floor(Math.random() * (max - min + 1) + min);
  }

  // StateMachineを実行するEventBridge Schedulerを設定する関数
  private createSchedulerForStateMachine(
    stateMachine: sfn.StateMachine,
    index: number
  ) {
    const target = new aws_scheduler_targets_alpha.StepFunctionsStartExecution(
      stateMachine,
      {}
    );

    new aws_scheduler_alpha.Schedule(this, `Schedule-StateMachine${index}`, {
      scheduleName: `invoke-StateMachine${index}-schedule`,
      schedule: aws_scheduler_alpha.ScheduleExpression.cron({
        minute: "0",
        hour: "0",
        day: "*",
        month: "*",
        year: "*",
        timeZone: TimeZone.ASIA_TOKYO,
      }),
      target,
    });
  }
}

import { StepFunctions, Athena } from "aws-sdk";
import { APIGatewayProxyResult } from "aws-lambda";
import { format, addHours } from "date-fns";

// AWS SDKのクライアントを初期化
const sfn = new StepFunctions();
const athena = new Athena();

// 環境変数から必要な値を取得
const BUCKET_NAME = process.env.BUCKET_NAME;
const ATHENA_WORKGROUP = process.env.ATHENA_WORKGROUP;
const DATABASE_NAME = process.env.DATABASE_NAME;
const TABLE_NAME = process.env.TABLE_NAME;

// StateMachine実行結果の型定義
interface ExecutionResult {
  StateMachineName: string;
  StartTime: string;
  Date: string;
  EndTime: string;
  Duration: number;
  Status: string;
}

// Lambda関数のメインハンドラー
export const handler = async (event: any): Promise<APIGatewayProxyResult> => {
  try {
    // 全StateMachineの実行結果を取得
    const results: ExecutionResult[] = await getAllExecutions();

    // 実行結果がない場合は処理をスキップ
    if (results.length === 0)
      return {
        statusCode: 200,
        body: JSON.stringify("Skip no data"),
      };

    // Athenaクエリを生成
    const query = buildAthenaQuery(results);
    // Athenaクエリを実行
    await executeAthenaQuery(query);

    // 成功レスポンスを返す
    return {
      statusCode: 200,
      body: JSON.stringify(
        "Daily State Machines execution summary inserted into Athena table"
      ),
    };
  } catch (error) {
    // エラーハンドリング
    console.error("Error:", error);
    return {
      statusCode: 500,
      body: JSON.stringify("An error occurred while processing the request"),
    };
  }
};

// 全StateMachineの実行結果を取得する関数
async function getAllExecutions(): Promise<ExecutionResult[]> {
  const stateMachines = await getStateMachines();
  let allResults: ExecutionResult[] = [];

  for (const stateMachineArn of stateMachines) {
    const executions = await listExecutions(stateMachineArn);
    allResults = allResults.concat(executions);
  }

  return allResults;
}

// 全StateMachineのARNを取得する関数
async function getStateMachines(): Promise<string[]> {
  const response = await sfn.listStateMachines().promise();
  return (
    response.stateMachines?.map(
      (sm: StepFunctions.StateMachineListItem) => sm.stateMachineArn ?? ""
    ) || []
  );
}

// UTCからJSTに変換する関数
function toJST(date: Date): Date {
  return addHours(date, 9);
}

// 日付をフォーマットする関数
function formatDate(date: Date): string {
  return format(toJST(date), "yyyy-MM-dd HH:mm:ss");
}

// 特定のStateMachineの実行結果を取得する関数
async function listExecutions(
  stateMachineArn: string
): Promise<ExecutionResult[]> {
  const results: ExecutionResult[] = [];
  const twentyFourHoursAgo = new Date(Date.now() - 24 * 60 * 60 * 1000);
  let nextToken: string | undefined;

  do {
    // StateMachineの実行履歴を取得
    const response = await sfn
      .listExecutions({
        stateMachineArn,
        maxResults: 100,
        nextToken,
      })
      .promise();

    if (response.executions) {
      for (const execution of response.executions) {
        // 24時間以内の完了した実行のみを対象とする
        if (
          execution.startDate >= twentyFourHoursAgo &&
          execution.status !== "RUNNING"
        ) {
          const endTime = execution.stopDate ? execution.stopDate : new Date();
          results.push({
            StateMachineName: stateMachineArn.split(":").pop() || "",
            StartTime: formatDate(execution.startDate),
            Date: format(toJST(execution.startDate), "yyyy-MM-dd"),
            Duration: (endTime.getTime() - execution.startDate.getTime()) / 1000,
            EndTime: formatDate(endTime),
            Status: execution.status,
          });
        } else if (execution.startDate < twentyFourHoursAgo) {
          // 24時間以上前の実行は対象外なのでループを終了
          nextToken = undefined;
          break;
        }
      }
    }

    nextToken = response.nextToken;
  } while (nextToken);

  return results;
}

// Athenaクエリを生成する関数
function buildAthenaQuery(results: ExecutionResult[]): string {
  const tableName = TABLE_NAME;
  const databaseName = DATABASE_NAME;

  // INSERTクエリの基本構造を作成
  let query = `
    INSERT INTO ${databaseName}.${tableName}
    SELECT 
      name,
      date,
      starttime,
      endtime,
      duration,
      status
    FROM (
  `;

  // 各実行結果をUNION ALLで結合
  for (let i = 0; i < results.length; i++) {
    const result = results[i];
    query += `
      SELECT
        '${result.StateMachineName}' as name,
        '${result.Date}' as date,
        '${result.StartTime}' as starttime,
        '${result.EndTime}' as endtime,
        ${result.Duration} as duration,
        '${result.Status}' as status
    `;

    if (i < results.length - 1) {
      query += "UNION ALL";
    }
  }

  query += ")";

  console.log(query);

  return query;
}

// Athenaクエリを実行する関数
async function executeAthenaQuery(query: string): Promise<void> {
  const params: Athena.StartQueryExecutionInput = {
    QueryString: query,
    QueryExecutionContext: {
      Database: DATABASE_NAME,
    },
    ResultConfiguration: {
      OutputLocation: `s3://${BUCKET_NAME}/athena-results/`,
    },
    WorkGroup: ATHENA_WORKGROUP,
  };

  // クエリを非同期で実行
  await athena.startQueryExecution(params).promise();
}

cdk deploy

select *
from statemachines_metrics